import numpy as np
from scipy.sparse import lil_matrix, csr_matrix
from scipy.sparse.linalg import spsolve
import matplotlib.pyplot as plt
import scienceplots
plt.style.use(['science'])

class FEM2D_Poisson:
    def __init__(self, Lx=1, Ly=1, nx=20, ny=20):
        self.Lx, self.Ly = Lx, Ly
        self.nx, self.ny = nx, ny
        self.hx = Lx / nx
        self.hy = Ly / ny
        self.nodes, self.elements = self.generate_mesh()

    def generate_mesh(self):
        x = np.linspace(0, self.Lx, self.nx + 1)
        y = np.linspace(0, self.Ly, self.ny + 1)
        nodes = np.array(np.meshgrid(x, y)).T.reshape(-1, 2)

        elements = []
        for i in range(self.ny):
            for j in range(self.nx):
                n0 = i * (self.nx + 1) + j
                n1 = n0 + 1
                n2 = (i + 1) * (self.nx + 1) + j + 1
                n3 = (i + 1) * (self.nx + 1) + j
                elements.append([n0, n1, n2, n3])

        return nodes, np.array(elements)

    def source_term(self, x, y):
        return -2

    def assemble_system(self):
        n_nodes = len(self.nodes)
        K = lil_matrix((n_nodes, n_nodes))
        F = np.zeros(n_nodes)

        gauss_points = np.array([[-1, -1], [1, -1], [1, 1], [-1, 1]]) * 0.57735
        weights = [1, 1, 1, 1]

        for elem in self.elements:
            node_coords = self.nodes[elem]
            Ke = np.zeros((4, 4))
            Fe = np.zeros(4)

            for gp, w in zip(gauss_points, weights):
                xi, eta = gp
                N = 0.25 * np.array([(1 - xi) * (1 - eta),
                                     (1 + xi) * (1 - eta),
                                     (1 + xi) * (1 + eta),
                                     (1 - xi) * (1 + eta)])
                dN_dxi = 0.25 * np.array([-(1 - eta), (1 - eta), (1 + eta), -(1 + eta)])
                dN_deta = 0.25 * np.array([-(1 - xi), -(1 + xi), (1 + xi), (1 - xi)])

                J = np.dot(np.array([dN_dxi, dN_deta]), node_coords)
                detJ = np.linalg.det(J)
                invJ = np.linalg.inv(J)

                dN_dx = invJ[0, 0] * dN_dxi + invJ[0, 1] * dN_deta
                dN_dy = invJ[1, 0] * dN_dxi + invJ[1, 1] * dN_deta

                Ke += (np.outer(dN_dx, dN_dx) + np.outer(dN_dy, dN_dy)) * detJ * w

                x_phys = np.dot(N, node_coords[:, 0])
                y_phys = np.dot(N, node_coords[:, 1])
                Fe += N * self.source_term(x_phys, y_phys) * detJ * w

            for i in range(4):
                for j in range(4):
                    K[elem[i], elem[j]] += Ke[i, j]
                F[elem[i]] += Fe[i]

        return K.tocsr(), F

    def apply_dirichlet_bc(self, K, F):
        tol = 1e-6
        boundary_nodes = []
        boundary_nodes += np.where(self.nodes[:, 0] < tol)[0].tolist()
        boundary_nodes += np.where(self.nodes[:, 0] > self.Lx - tol)[0].tolist()
        boundary_nodes += np.where(self.nodes[:, 1] < tol)[0].tolist()
        boundary_nodes += np.where(self.nodes[:, 1] > self.Ly - tol)[0].tolist()

        for node in boundary_nodes:
            K[node, :] = 0
            K[:, node] = 0
            K[node, node] = 1
            F[node] = 0

        return K, F

    def solve(self):
        K, F = self.assemble_system()
        K, F = self.apply_dirichlet_bc(K, F)
        u = spsolve(K, F)
        return u

    def exact_solution(self, x, y, terms=10):
        u_sum = np.zeros_like(x)
        for n in range(1, terms, 2):
            for m in range(1, terms, 2):
                coefficient = -32 / (n * m * np.pi ** 2 * ((n * np.pi / self.Lx) ** 2 + (m * np.pi / self.Ly) ** 2))
                u_sum += coefficient * np.sin(n * np.pi * x / self.Lx) * np.sin(m * np.pi * y / self.Ly)
        return u_sum

    def compare_with_exact(self, u_numerical, terms=10):
        X = self.nodes[:, 0].reshape(self.ny + 1, self.nx + 1)
        Y = self.nodes[:, 1].reshape(self.ny + 1, self.nx + 1)
        U_numerical = u_numerical.reshape(self.ny + 1, self.nx + 1)
        U_exact = self.exact_solution(X, Y, terms)
        error = np.abs(U_numerical - U_exact)


        # ==================== 数值解（二维伪彩色图） ====================
        fig1 = plt.figure(figsize=(4, 3))
        ax1 = fig1.add_subplot(111)
        mesh1 = ax1.pcolormesh(X, Y, U_numerical, cmap='viridis', shading='auto')
        fig1.colorbar(mesh1, ax=ax1, label='$u(x, y)$')
        ax1.set_title('Numerical Solution')
        ax1.set_xlabel('$x$')
        ax1.set_ylabel('$y$')
        plt.tight_layout()
        plt.show()
        fig1.savefig('./numerical_solution1.png', dpi=300)

        # ==================== 解析解（二维等高线图） ====================
        fig2 = plt.figure(figsize=(4, 3))
        ax2 = fig2.add_subplot(111)
        contour = ax2.contourf(X, Y, U_exact, levels=20, cmap='viridis', shading='auto')
        fig2.colorbar(contour, ax=ax2, label='$u(x, y)$')
        ax2.set_title('Exact Solution')
        ax2.set_xlabel('$x$')
        ax2.set_ylabel('$y$')
        plt.tight_layout()
        plt.show()
        fig2.savefig('./exact_solution1.png', dpi=300)

        # ==================== 误差图（二维伪彩色图） ====================
        fig3 = plt.figure(figsize=(4, 3))
        ax3 = fig3.add_subplot(111)
        mesh3 = ax3.pcolormesh(X, Y, error, cmap='coolwarm', shading='auto')
        fig3.colorbar(mesh3, ax=ax3, label='Error')
        # ax3.set_title('Error')
        ax3.set_xlabel('$x$')
        ax3.set_ylabel('$y$')
        plt.tight_layout()
        plt.show()
        fig3.savefig('./error1.png', dpi=300)

        print("Maximum error:", np.max(error))
        print("L2 norm of error:", np.linalg.norm(error))

if __name__ == "__main__":
    fem = FEM2D_Poisson(nx=30, ny=30)
    u_numerical = fem.solve()
    fem.compare_with_exact(u_numerical)